#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2018 crane <crane@crane-pc>
#
# Distributed under terms of the MIT license.


"""
线性回归, 计算机求解
"""

import random
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D # 绘制3D坐标的函数


class Regression:
    def __init__(self):
        pass

    def regress(self, xs, ys):
        self.xs = xs
        self.ys = ys
        self.n  = len(xs)

        self._scatter()
        self._best_fit_line()
        self._plot_residual()
        self._plot_k_b_3d()     # k,b关于二乘残差的二元函数图(3D图)
        self._corelation_efficient()

        plt.show()

    def _scatter(self):
        plt.scatter(self.xs, self.ys, c='r')

    def _plot_line(self):
        # xs, ys = self.xs, self.ys
        xmin, xmax = min(self.xs)-1, max(self.xs) + 1
        # xmin, xmax = 0, max(self.xs) + 1

        f = lambda x : self.k * x + self.b
        ymin, ymax = f(xmin), f(xmax)

        xs = (xmin, xmax)
        ys = (ymin, ymax)
        plt.plot(xs, ys)

    def _best_fit_line(self):
        xs, ys = self.xs, self.ys

        x_mean = sum(xs) / self.n
        y_mean = sum(ys) / self.n

        x_square = [x*x for x in xs]
        x_square_mean = sum(x_square) / self.n

        y_square = [y*y for y in ys]
        y_square_mean = sum(y_square) / self.n

        xy = list( map(lambda x,y: x * y, xs, ys) )
        xy_mean = sum(xy) / len(xy)

        self.x_mean, self.y_mean, self.x_square_mean, self.y_square_mean, self.xy_mean = \
                x_mean, y_mean, x_square_mean, y_square_mean, xy_mean


        point1 = (x_mean, y_mean)
        point2 = (x_square_mean/x_mean, xy_mean/x_mean)

        self.k, self.b = k_b(point1, point2)
        print('=================== (k, b) (%s, %s)=====================' %(self.k, self.b))
        self._plot_line()

    def _plot_residual(self):
        ''' 画出残差 '''
        for i, x in enumerate(self.xs):
            y = self.ys[i]
            fit_y = self.k * x + self.b
            plt.plot([x, x], [y, fit_y], c='y')

    def _plot_k_b_3d(self):
        x_mean, y_mean, x_square_mean, y_square_mean, xy_mean = \
                self.x_mean, self.y_mean, self.x_square_mean, self.y_square_mean, self.xy_mean
        # k = range(1, 6)
        # b = range(1, 6)
        k = np.arange(1, 6, 0.2)
        b = np.arange(1, 6, 0.2)
        K, B = np.meshgrid(k,b)

        sub = np.subtract
        add = np.add
        mul = np.multiply

        def f(k, b):
            # y_square_mean - 2*xy_mean*k - 2*y_mean*b + k*k*x_square_mean + 2*k*b*x_mean + b*b
            item1 = y_square_mean
            item2 = mul(k, 2*xy_mean)
            item3 = mul(b, 2*y_mean)
            item4 = mul(mul(k, k), x_square_mean)
            item5 = mul(mul(k, b), 2*x_mean)
            item6 = mul(b, b)

            part1 = sub(sub(item1, item2), item3)
            part2 = add(add(item4, item5), item6)
            return add(part1, part2)
        Z = f(K, B)

        # print(Z)
        fig2=plt.figure()                                                    # 创建一个绘图对象
        ax=Axes3D(fig2)                                                      # 用这个绘图对象创建一个Axes对象(有3D坐标)
        ax.plot_surface(K, B, Z, rstride=1, cstride=1, cmap=plt.cm.coolwarm) # 用取样点(x,y,z)去构建曲面
        ax.set_xlabel('K label', color='r')
        ax.set_ylabel('B label', color='g')
        ax.set_zlabel('Z label', color='b')                                  # 给三个坐标轴注明

        # ax.scatter([self.k], [self.b], add(f([self.k], [self.b]), 1), c='red', marker='o', linewidths=2)
        min_x, min_y, min_z = self.k, self.b, f(self.k, self.b)
        # print(k)
        # print('v %s' % v)
        length = 20000
        ax.plot([min_x, min_x], [min_y, min_y], [min_z-length, min_z+length])

        # TODO 在3D图最小点处, 点一个红点

    def _fit_value(self, x):
        return self.k * x + self.b

    def _residual(self, x, y):
        return y - self._fit_value(x)

    def _corelation_efficient(self):
        # 不受回归方程影响的波动
        se_line = map(
            lambda x, y : self._residual(x, y) ** 2,
            self.xs, self.ys
        )

        # 总的影响
        se_average  = map(
            lambda y: (y - self.y_mean) ** 2,
            self.ys
        )

        correlation_efficient = 1 - ( sum(se_line) / sum(se_average) )
        print('=================== correlation_efficient %s =====================' % correlation_efficient)


def k_b(point1, point2):
    ''' 给定两点: 求斜率和截距 '''
    x1, y1 = point1
    x2, y2 = point2

    k = (y2 - y1) / (x2 - x1)
    b = y1 - k * x1
    return k,b

def random_sign():
    return random.choice([1, -1])

def generate_input(k=2, b=4, n=100):
    ''' k: 斜率, b: 截距'''
    xs = list(range(n))

    def y(x):
        # 为y添加[0, 1]的干扰, 干扰随机浮动于正负之间
        # interference = random.random()
        interference = random.random() * random_sign() * 4
        yvalue = k*x + b + interference
        return yvalue

    ys = list( map(y, xs) )
    print(xs, ys)
    return xs, ys


def test():
    # xs, ys = generate_input()
    xs, ys = [
        [1, 2, 3, 4],
        [1, 3, 8, 9],
    ]

    r = Regression()
    r.regress(xs, ys)

def main():
    print("start main")
    test()

if __name__ == "__main__":
    main()
